// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <nebula/compute/kernels/codegen_internal.h>

namespace nebula::compute::internal {

    template<typename Type, typename Enable = void>
    struct CopyDataUtils {
    };

    template<>
    struct CopyDataUtils<BooleanType> {
        static void CopyData(const DataType &, const Scalar &in, const int64_t in_offset,
                             uint8_t *out, const int64_t out_offset, const int64_t length) {
            bit_util::SetBitsTo(
                    out, out_offset, length,
                    in.is_valid ? turbo::checked_cast<const BooleanScalar &>(in).value : false);
        }

        static void CopyData(const DataType &, const uint8_t *in, const int64_t in_offset,
                             uint8_t *out, const int64_t out_offset, const int64_t length) {
            nebula::internal::CopyBitmap(in, in_offset, length, out, out_offset);
        }

        static void CopyData(const DataType &, const ArraySpan &in, const int64_t in_offset,
                             uint8_t *out, const int64_t out_offset, const int64_t length) {
            const auto in_arr = in.get_values<uint8_t>(1, /*absolute_offset=*/0);
            CopyData(*in.type, in_arr, in_offset, out, out_offset, length);
        }
    };

    template<>
    struct CopyDataUtils<FixedSizeBinaryType> {
        static void CopyData(const DataType &ty, const Scalar &in, const int64_t in_offset,
                             uint8_t *out, const int64_t out_offset, const int64_t length) {
            const int32_t width = ty.byte_width();
            uint8_t *begin = out + (width * out_offset);
            const auto &scalar = turbo::checked_cast<const nebula::internal::PrimitiveScalarBase &>(in);
            // Null scalar may have null value buffer
            if (!scalar.is_valid) {
                std::memset(begin, 0x00, width * length);
            } else {
                const std::string_view buffer = scalar.view();
                        DKCHECK_GE(buffer.size(), static_cast<size_t>(width));
                for (int i = 0; i < length; i++) {
                    std::memcpy(begin, buffer.data(), width);
                    begin += width;
                }
            }
        }

        static void CopyData(const DataType &ty, const uint8_t *in, const int64_t in_offset,
                             uint8_t *out, const int64_t out_offset, const int64_t length) {
            const int32_t width = ty.byte_width();
            uint8_t *begin = out + (width * out_offset);
            std::memcpy(begin, in + in_offset * width, length * width);
        }

        static void CopyData(const DataType &ty, const ArraySpan &in, const int64_t in_offset,
                             uint8_t *out, const int64_t out_offset, const int64_t length) {
            const int32_t width = ty.byte_width();
            const auto in_arr = in.get_values<uint8_t>(1, in.offset * width);
            CopyData(ty, in_arr, in_offset, out, out_offset, length);
        }
    };

    template<typename Type>
    struct CopyDataUtils<
            Type, enable_if_t<is_number_type<Type>::value || is_interval_type<Type>::value>> {
        using CType = typename TypeTraits<Type>::CType;

        static void CopyData(const DataType &, const Scalar &in, const int64_t in_offset,
                             uint8_t *out, const int64_t out_offset, const int64_t length) {
            CType *begin = reinterpret_cast<CType *>(out) + out_offset;
            CType *end = begin + length;
            std::fill(begin, end, UnboxScalar<Type>::Unbox(in));
        }

        static void CopyData(const DataType &, const uint8_t *in, const int64_t in_offset,
                             uint8_t *out, const int64_t out_offset, const int64_t length) {
            std::memcpy(out + out_offset * sizeof(CType), in + in_offset * sizeof(CType),
                        length * sizeof(CType));
        }

        static void CopyData(const DataType &, const ArraySpan &in, const int64_t in_offset,
                             uint8_t *out, const int64_t out_offset, const int64_t length) {
            const auto in_arr = in.get_values<uint8_t>(1, in.offset * sizeof(CType));
            CopyData(*in.type, in_arr, in_offset, out, out_offset, length);
        }
    };

}  // namespace nebula::compute::internal
